// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: jaewon

#include "coin2/base/request/session_pool.h"

#include <cmath>
#include <cstdlib>
#include <functional>
#include <regex>
#include <string>
#include <utility>

#include <boost/asio/dispatch.hpp>
#include <boost/asio/spawn.hpp>
#include <boost/beast/version.hpp>
#include <boost/exception/diagnostic_information.hpp>
#include <glog/logging.h>

#include "coin2/base/log.h"
#include "coin2/base/request/root_certificates.h"

// #define ENABLE_HTTP_PIPELINING

namespace coin2::base::request {
namespace detail {

static const auto kMaxSessionRttCount = 30LL;

static RequestKey ParserRequest(const std::string& method, const std::string& target) {
  RequestKey key;
  key.first = method;
  auto size = target.size();
  for (auto i = 0; i < size; ++i) {
    if (target[i] == '?') {
      key.second = target.substr(0, i);
      return key;
    }
  }
  key.second = target;
  return key;
}

bool operator<(const SessionRtt& lhs, const SessionRtt& rhs) { return lhs.rtt < rhs.rtt; }

SingleHostSessionPool::SingleHostSessionPool(
    boost::asio::io_context* io_ctx,
    ssl::context* ssl_ctx,
    const std::string& host,
    int port,
    const std::string& bind_ip,
    int max_sessions)
    : io_ctx_(io_ctx),
      ssl_ctx_(ssl_ctx),
      host_(host),
      port_(port),
      bind_ip_(bind_ip),
      max_sessions_(max_sessions) {
  req_ctx_pool_ = std::make_shared<detail::SimpleObjectPool<Session::RequestContext>>();
  is_prex_ = std::regex_search(host_, std::regex("(aqx|prex)"));
}

SingleHostSessionPool::~SingleHostSessionPool() { AsyncShutdown(); }

bool SingleHostSessionPool::is_blocking() const {
  if (max_sessions_ > 0 && work_sessions_.size() + free_sessions_.size() == max_sessions_) {
    if (free_sessions_.empty()) {
      return true;
    }
  }
  return false;
}

auto SingleHostSessionPool::FindMinOrMaxRttSession(bool is_max)
    -> std::unordered_set<std::shared_ptr<Session>>::iterator {
  auto it = free_sessions_.begin();
  if (!is_max) {
    while (!rtt_tree_.empty()) {
      auto now = rtt_tree_.begin();
      auto se = now->session.lock();
      rtt_tree_.erase(now);
      if (!se) {
        continue;
      }
      if (free_sessions_.find(se) != free_sessions_.end()) {
        return free_sessions_.find(se);
      }
    }
  } else {
    while (!rtt_tree_.empty()) {
      auto now = --rtt_tree_.end();
      auto se = now->session.lock();
      rtt_tree_.erase(now);
      if (!se) {
        continue;
      }
      if (free_sessions_.find(se) != free_sessions_.end()) {
        return free_sessions_.find(se);
      }
    }
  }
  return it;
}

#ifdef DISABLE_CHOOSE_FASTEST_SESSION

void SingleHostSessionPool::OnSessionReady(
    std::weak_ptr<Session> session,
    int64_t rtt,
    const std::string& req_method,
    const std::string& req_target,
    int32_t http_code) {
  if (is_prex_) {
    StatsRequestInfo(req_method, req_target, http_code);
  }
  rtts_.Calculate(rtt / 1000000);

  auto se = session.lock();
  if (se == nullptr) {
    return;
  }

  if (max_sessions_ > 0 && work_sessions_.size() + free_sessions_.size() > max_sessions_) {
    DeferDeleteSession(session, true);
    return;
  }

  work_sessions_.erase(se);
  timeout_sessions_.erase(se);
  free_sessions_.insert(std::move(se));
}

#else

void SingleHostSessionPool::OnSessionReady(
    std::weak_ptr<Session> session,
    int64_t rtt,
    const std::string& req_method,
    const std::string& req_target,
    int32_t http_code) {
  if (is_prex_) {
    StatsRequestInfo(req_method, req_target, http_code);
  }
  rtts_.Calculate(rtt / 1000000);

  if (http_code >= 500) {  // TODO(daniel): can we treat this session is not available any more?
    DeferDeleteSession(session, true);
    return;
  }

  auto se = session.lock();
  if (se == nullptr) {
    return;
  }
  auto timestamp = se->GetLastActiveTs();
  auto request_priority_weight = se->GetRequestPriorityWeight();
  if (rtt > 0 && request_priority_weight == RequestPriorityWeight::kHigh /* SUBMIT/CANCEL */) {
    auto tmp = rtt + rtt;
    if (0 == sessions_rtt_list_.count(se)) {
      sessions_rtt_list_[se] = {rtt};
    } else {
      int cnt = sessions_rtt_list_[se].size();
      if (sessions_rtt_list_[se].size() < kMaxSessionRttCount) {
        int i = cnt;
        sessions_rtt_list_[se].push_back(rtt);
        while (i > 0 && rtt < sessions_rtt_list_[se][i - 1]) {
          sessions_rtt_list_[se][i] = sessions_rtt_list_[se][i - 1];
          i--;
        }
        sessions_rtt_list_[se][i] = rtt;
      }
      tmp = (sessions_rtt_list_[se][cnt >> 1] + sessions_rtt_list_[se][(cnt - 1) >> 1]);
    }
    rtt_tree_.insert({tmp, se});
  } else {
    if (0 == sessions_rtt_list_.count(se)) {
      rtt_tree_.insert({0, se});
    }
  }

  work_sessions_.erase(se);
  timeout_sessions_.erase(se);
  free_sessions_.insert(std::move(se));

  UpdateRequestWindow(timestamp, request_priority_weight);
  // release session if exceed the max limitation
  if (max_sessions_ > 0 && work_sessions_.size() + free_sessions_.size() > max_sessions_) {
    auto it = FindMinOrMaxRttSession(true);
    auto session = *it;
    free_sessions_.erase(it);
    DeferDeleteSession(session, false);
  }
}

#endif

void SingleHostSessionPool::OnSessionError(std::weak_ptr<Session> session) {
  rtts_.ErrNum();
  // req_status error
  DeferDeleteSession(session, true);
}

void SingleHostSessionPool::StatsRequestInfo(
    const std::string& req_method,
    const std::string& req_target,
    int32_t code) {
  auto key = ParserRequest(req_method, req_target);
  if (req_status_.find(key) == req_status_.end()) {
    RequestInfo info;
    info.AddInfo(code);
    req_status_.emplace(std::make_pair(key, std::move(info)));
  } else {
    req_status_[key].AddInfo(code);
  }
}

void SingleHostSessionPool::UpdateRequestWindow(
    int64_t timestamp,
    RequestPriorityWeight priority_weight) {
  if (RequestPriorityWeight::kHigh == priority_weight) {
    free_session_add_records_list_.push_back(timestamp);
  } else {
    // TODO: RequestPriorityWeight::kCritical == priority_weight
    working_query_num_ -= 1;
  }
}

void SingleHostSessionPool::DeferDeleteSession(
    std::weak_ptr<Session> session,
    bool request_window_need_updating) {
  auto se = session.lock();
  if (se == nullptr) {
    return;
  }

  if (request_window_need_updating) {
    auto timestamp =
        RequestPriorityWeight::kHigh == se->GetRequestPriorityWeight() ? GetCurrentTimestamp() : 0;
    UpdateRequestWindow(timestamp, se->GetRequestPriorityWeight());
  }

  // If this is a work session
  work_sessions_.erase(se);
  timeout_sessions_.erase(se);
  // If this is an free session
  free_sessions_.erase(se);
  if (!del_sessions_.count(se)) {
    se->AsyncShutdown();
    del_sessions_.insert(std::move(se));
  }
}

void SingleHostSessionPool::DeleteSession() {
  auto iter = del_sessions_.begin();
  while (iter != del_sessions_.end()) {
    auto conn = (*iter)->Connection().lock();
    if (conn && conn->in_sending()) {
      ++iter;
      continue;
    }
    sessions_rtt_list_.erase(*iter);
    iter = del_sessions_.erase(iter);
  }
}

int SingleHostSessionPool::PoissonDistributionSampling() {
  // Poisson distribution
  double avg_del_num =
      (free_session_del_records_list_.size() - 1) * 1.0e6 /
      (free_session_del_records_list_.back() - free_session_del_records_list_.front());
  double avg_add_num =
      (free_session_add_records_list_.size() - 1) * 1.0e6 /
      (free_session_add_records_list_.back() - free_session_add_records_list_.front());
  std::poisson_distribution<int> add_distribution(avg_add_num * 100);
  std::poisson_distribution<int> del_distribution(avg_del_num * 100);
  // sampling
  int sampling_result = del_distribution(generator_) - add_distribution(generator_);
  return sampling_result;
}

bool SingleHostSessionPool::WhetherNeedToPreCreateSession() {
  if (free_session_del_records_list_.size() >= session_window_size_ &&
      free_session_add_records_list_.size() >= session_window_size_ ) {
      int k = free_sessions_.size();
      int sampling_result = PoissonDistributionSampling();
      if (sampling_result >= k) {
        return true;
      }
  }
  return false;
}

bool SingleHostSessionPool::CheckIfFreeSessionsAreEnoughForQuery(
    RequestPriorityWeight priority_weight) {
  while (free_session_del_records_list_.size() > session_window_size_) {
    free_session_del_records_list_.pop_front();
  }
  while (free_session_add_records_list_.size() > session_window_size_) {
    free_session_add_records_list_.pop_front();
  }
  if (RequestPriorityWeight::kHigh == priority_weight) {
    free_session_del_records_list_.push_back(GetCurrentTimestamp());
  } else {
    if (RequestPriorityWeight::kCritical == priority_weight) {
      // always assign session
    } else {
      if (free_session_del_records_list_.size() >= session_window_size_ &&
          free_session_add_records_list_.size() >= session_window_size_ &&
          working_query_num_ >= max_sessions_ * min_proportion_of_query_) {
        int k = max_sessions_ - work_sessions_.size();
        int sampling_result = PoissonDistributionSampling();
        if (sampling_result >= k) {
          return false;
        }
      }
    }
    working_query_num_ += 1;
  }
  return true;
}

#ifdef DISABLE_CHOOSE_FASTEST_SESSION

std::shared_ptr<Session> SingleHostSessionPool::TakeSession(RequestPriorityWeight priority_weight) {
  // Connection limit
  if (max_sessions_ > 0 && work_sessions_.size() + free_sessions_.size() >= max_sessions_) {
    if (free_sessions_.empty()) {
      return nullptr;
    }
  }

  if (!CheckIfFreeSessionsAreEnoughForQuery(priority_weight)) {
    return nullptr;
  }

  if (free_sessions_.empty() || (priority_weight != RequestPriorityWeight::kHigh &&
                                 rand() % (free_sessions_.size() + 1) == 0)) {
    auto session = std::make_shared<Session>(
        io_ctx_,
        this,
        req_ctx_pool_,
        ssl_ctx_,
        host_,
        port_,
        bind_ip_,
        http_proxy_);
    return session;
  }
  auto it = free_sessions_.begin();
  auto session = *it;
  free_sessions_.erase(it);
  return session;
}

#else

std::shared_ptr<Session> SingleHostSessionPool::TakeSession(RequestPriorityWeight priority_weight) {
  // Connection limit
  if (max_sessions_ > 0 && work_sessions_.size() + free_sessions_.size() >= max_sessions_) {
    if (free_sessions_.empty()) {
      return nullptr;
    }
  }

  if (!CheckIfFreeSessionsAreEnoughForQuery(priority_weight)) {
    return nullptr;
  }

  if (free_sessions_.empty() || (priority_weight != RequestPriorityWeight::kHigh &&
                                 WhetherNeedToPreCreateSession())) {
    auto session = std::make_shared<Session>(
        io_ctx_,
        this,
        req_ctx_pool_,
        ssl_ctx_,
        host_,
        port_,
        bind_ip_,
        http_proxy_);
    return session;
  }
  auto it = FindMinOrMaxRttSession(priority_weight != RequestPriorityWeight::kHigh);
  auto session = *it;
  CHECK(session) << "#og_team, please check that failed to get http session";
  free_sessions_.erase(it);
  return session;
}

#endif

void SingleHostSessionPool::AsyncRequest(
    Session::RequestType req,
    Session::CallbackType callback,
    RequestPriorityWeight priority_weight) {
  auto session = TakeSession(priority_weight);
  if (session == nullptr) {
    LOG(WARNING) << "No available connection!";
    io_ctx_->post([req = std::move(req), callback = std::move(callback)]() {
      auto error_code =
          boost::system::errc::make_error_code(boost::system::errc::too_many_files_open);
      auto ts = GetCurrentTimestamp();
      SessionResponse response;
      response.error_code = error_code;
      response.raw_request = &req;
      response.request_ts = ts;
      response.response_ts = ts;
      callback(error_code, &response);
    });
    return;
  }
  rtts_.ReqNum();
  session->SetRequestPriorityWeight(priority_weight);
  session->AsyncRequest(std::move(req), std::move(callback));
  work_sessions_.insert(session);
}

void SingleHostSessionPool::AsyncShutdown() {
  for (auto& s : work_sessions_) {
    s->AsyncShutdown();
  }
  for (auto& s : free_sessions_) {
    s->AsyncShutdown();
  }
}

void SingleHostSessionPool::ProcessTimeout(int64_t timeout) {
  auto now = GetCurrentTimestamp();
  auto timeout_start_ts = now - timeout;
  for (auto& session : work_sessions_) {
    auto ret = session->ProcessTimeout(timeout_start_ts);
    if (!ret) {
      continue;
    }

    timeout_sessions_.insert(session);
  }

  ShutdownTimeoutSession(now);

  // TODO(daniel): handle timeout session is active for long time
  if (timeout_sessions_.size() > 50) {
    SPDLOG_WARN(
        "too many timeout sessions({}-{}-{}) are not recycled yet",
        timeout_sessions_.size(),
        work_sessions_.size(),
        free_sessions_.size());
  }

  DeleteSession();
}

void SingleHostSessionPool::ShutdownTimeoutSession(int64_t now) {
  if (now - last_shutdown_timeout_session_ts_ <= check_shutdown_timeout_session_period_ns_) {
    return;
  }
  last_shutdown_timeout_session_ts_ = now;

  std::vector<std::shared_ptr<Session>> need_delete_sessions;

  for (auto& session : timeout_sessions_) {
    auto last_request_ts = session->GetLastRequestTs();
    if (now - last_request_ts <= shutdown_timeout_session_ns_) {
      continue;
    }
    need_delete_sessions.emplace_back(session);
  }

  for (const auto& session : need_delete_sessions) {
    LOG(ERROR) << "Shutdown timeout session, time from last request ns:"
               << now - session->GetLastRequestTs() << " request:\n"
               << session->GetRequest();
    DeferDeleteSession(session, true);
  }
}

void SingleHostSessionPool::OnTimer(int64_t timeout) {
  rtts_.ReportSessionSize(work_sessions_.size(), free_sessions_.size());
  ProcessTimeout(timeout);
}

int SingleHostSessionPool::TotalNumReqs() { return work_sessions_.size(); }

RttInfo SingleHostSessionPool::GetRttInfo() {
  RttInfo rtts = rtts_;
  rtts_.Reset();
  return rtts;
}

RequestStatus SingleHostSessionPool::GetRequestInfo() {
  RequestStatus info;
  info.swap(req_status_);
  return info;
}

void DoQuery(
    std::string url_root,
    std::string url_path,
    std::string* content,
    boost::asio::io_context* io_ctx,
    boost::asio::yield_context yield) {
  try {
    std::shared_ptr<request::SessionPool> pool(new request::SessionPool(io_ctx));
    boost::system::error_code ec;
    boost::asio::steady_timer timer(*io_ctx);
    pool->SetTimeout(120 * 1000000000L);
    auto res = pool->CoroRequest(url_root, 443, url_path, yield[ec]);
    if (!ec) {
      *content = res->raw_response->body();
    } else {
      CHECK(false) << "Query failed. Error code: " << ec;
    }

    timer.expires_from_now(std::chrono::seconds(3));
    timer.async_wait(yield[ec]);

    pool->AsyncShutdown();
  } catch (...) {
    LOG(ERROR) << boost::current_exception_diagnostic_information();
    CHECK(false) << "Query failed.";
  }
}

}  // namespace detail

int64_t GetOverageFromEnvVar() {
  char const* overage = getenv("HTTP_CONN_DURATION");
  if (overage != nullptr) {
    LOG(INFO) << "Http connection overage is set as " << overage << " seconds";
    return std::stoi(std::string(overage)) * 1000000000L;
  }
  return std::numeric_limits<int64_t>::max();
}

SessionPool::SessionPool(boost::asio::io_context* io_ctx)
    : io_ctx_(io_ctx), max_sessions_per_host_(0), timer_(*io_ctx_) {
  ssl_ctx_ = GetDefaultSslContext();
}

SessionPool::SessionPool(boost::asio::io_context* io_ctx, int max_sessions_per_host)
    : io_ctx_(io_ctx), max_sessions_per_host_(max_sessions_per_host), timer_(*io_ctx_) {
  ssl_ctx_ = GetDefaultSslContext();
  CHECK_GE(max_sessions_per_host, 0);
}

SessionPool::SessionPool(
    boost::asio::io_context* io_ctx,
    ssl::context* ssl_ctx,
    int max_sessions_per_host)
    : io_ctx_(io_ctx),
      ssl_ctx_(ssl_ctx),
      max_sessions_per_host_(max_sessions_per_host),
      timer_(*io_ctx_) {
  CHECK_GE(max_sessions_per_host, 0);
}

SessionPool::~SessionPool() = default;

void SessionPool::SetTimeout(int64_t timeout_ns) { timeout_ns_ = timeout_ns; }

void SessionPool::StartTimer() {
  timer_enabled_ = true;
  overage_ns_ = GetOverageFromEnvVar();
  timer_.expires_after(std::chrono::seconds(1));
  timer_.async_wait(std::bind(&SessionPool::OnTimer, shared_from_this(), std::placeholders::_1));
  LOG(INFO) << "overage_sec = " << overage_ns_ / 1e9;
}

void SessionPool::OnTimer(const boost::system::error_code& ec) {
  if (ec || !timer_enabled_) return;

  if (timeout_ns_ != 0) {
    for (auto& pair : pools_) {
      pair.second->OnTimer(timeout_ns_);
    }
  }
  // disable this checking due to policy changed by huobi
  // if (overage_ns_ != 0) {
  //   for (auto& pair : pools_) {
  //     pair.second->ProcessOverage(overage_ns_);
  //   }
  // }

  timer_.expires_after(std::chrono::seconds(1));
  timer_.async_wait(std::bind(&SessionPool::OnTimer, shared_from_this(), std::placeholders::_1));
}

void SessionPool::AsyncRequest(
    const std::string& host,
    int port,
    const std::string& bind_ip,
    Session::RequestType req,
    Session::CallbackType callback,
    RequestPriorityWeight priority_weight,
    http::verb method) {
  if (!timer_enabled_) StartTimer();

  auto pool = GetSingleHostSessionPool(host, port, bind_ip);
  pool->AsyncRequest(std::move(req), callback, priority_weight);
}

void SessionPool::AsyncRequest(
    const std::string& host,
    int port,
    const std::string& bind_ip,
    std::string_view target,
    Session::CallbackType callback,
    http::verb method) {
  Session::RequestType req;

  const int version = 11;
  req.version(version);
  req.method(method);
  req.keep_alive(true);
  req.target(boost::beast::string_view(target.data(), target.size()));
  req.set(http::field::host, host);
  req.set(http::field::user_agent, BOOST_BEAST_VERSION_STRING);

  AsyncRequest(host, port, bind_ip, std::move(req), std::move(callback));
}

void SessionPool::AsyncRequest(
    const std::string& host,
    int port,
    const std::string& bind_ip,
    Session::RequestType req,
    Session::CallbackType callback,
    RequestPriorityWeight priority_weight) {
  AsyncRequest(host, port, bind_ip, std::move(req), callback, priority_weight, http::verb::get);
}

void SessionPool::AsyncShutdown() {
  timer_enabled_ = false;
  timer_.cancel();

  for (auto& pair : pools_) {
    pair.second->AsyncShutdown();
  }
}

// TODO(session pool is now used as single host session pool)
bool SessionPool::is_blocking() const {
  for (const auto& pool : pools_) {
    return pool.second->is_blocking();
  }
  return false;
}

detail::SingleHostSessionPool* SessionPool::GetSingleHostSessionPool(
    const std::string& host,
    int port,
    const std::string& bind_ip) {
  Key key(host, port, bind_ip);
  auto it = pools_.find(key);
  if (it != pools_.end()) {
    return it->second.get();
  }
  auto pool = std::make_unique<detail::SingleHostSessionPool>(
      io_ctx_,
      ssl_ctx_,
      host,
      port,
      bind_ip,
      max_sessions_per_host_);
  pool->SetHttpProxy(http_proxy_);
  auto ret = pool.get();
  pools_[key] = std::move(pool);
  return ret;
}

RttInfo SessionPool::GetRttInfo() {
  RttInfo rtt;
  for (const auto& pool : pools_) {
    rtt.Add(pool.second->GetRttInfo());
  }
  return rtt;
}

RequestStatus SessionPool::GetRequestInfo() {
  RequestStatus info;
  for (auto& pool : pools_) {
    auto temp = pool.second->GetRequestInfo();
    info.merge(temp);
    for (auto x : temp) {
      info[x.first] += x.second;
    }
  }
  return info;
}

std::string GetHttpsContent(const std::string& url_root, const std::string& url_path) {
  boost::asio::io_context io_ctx;
  std::string content;
  boost::asio::spawn(
      io_ctx,
      std::bind(detail::DoQuery, url_root, url_path, &content, &io_ctx, std::placeholders::_1));
  io_ctx.run();
  return content;
}
}  // namespace coin2::base::request
